/*
 * Copyright 1993-2010 NVIDIA Corporation.  All rights reserved.
 *
 * Please refer to the NVIDIA end user license agreement (EULA) associated
 * with this source code for terms and conditions that govern your use of
 * this software. Any use, reproduction, disclosure, or distribution of
 * this software and related documentation outside the terms of the EULA
 * is strictly prohibited.
 *
 */

///////////////////////////////////////////////////////////////////////////////
#include <cufft.h>
#include <math_constants.h>
#include "water_cu.h"

//#include <sm_11_atomic_functions.h>

//Round a / b to nearest higher integer value
int cuda_iDivUp(int a, int b)
{
    return (a + (b - 1)) / b;
}


// complex math functions
__device__																	    
float2 conjugate(float2 arg)
{
    return make_float2(arg.x, -arg.y);
}

__device__
float2 complex_exp(float arg)
{
    return make_float2(cosf(arg), sinf(arg));
}

__device__
float2 complex_add(float2 a, float2 b)
{
    return make_float2(a.x + b.x, a.y + b.y);
}

__device__
float2 complex_mult(float2 ab, float2 cd)
{
    return make_float2(ab.x * cd.x - ab.y * cd.y, ab.x * cd.y + ab.y * cd.x);
}

// generate wave heightfield at time t based on initial heightfield and dispersion relationship
__global__ void generateSpectrumKernel(float2* h0, float2 *ht, unsigned int width, unsigned int height, float t, float dk)
{
    unsigned int x = blockIdx.x*blockDim.x + threadIdx.x;
    unsigned int y = blockIdx.y*blockDim.y + threadIdx.y;
    unsigned int i = y*width+x;
    
    // calculate coordinates
    float2 k;
    k.x = x * dk;
    k.y = y * dk;
	if (x > width / 2)
	{
		k.x -= dk*width;
	}
	if (y > height / 2)
	{
		k.y -= dk*height;
	}

    // calculate dispersion w(k)
    float k_len = sqrtf(k.x*k.x + k.y*k.y);
    float w = sqrtf(9.8f * k_len);

	float2 h_tilda;
	if ((x < width) && (y < height)) {
		float2 h0_k = h0[i];
		//float2 h0_mk = h0[(((height-1)-y)*width)+x];

		//h_tilda = complex_add( complex_mult(h0_k, complex_exp(w * t)),
        //                          complex_mult(conjugate(h0_mk), complex_exp(-w * t)) );
        
        h_tilda = complex_mult(h0_k, complex_exp(-w * t));
	}

    // output frequency-space complex values
    if ((x < width) && (y < height)) {
        ht[i] = h_tilda;
    }
}


// generate slope by partial differences in spatial domain
__global__ void calculateSlopeKernel(float* h, float2 *slopeOut, unsigned int width, unsigned int height)
{
    unsigned int x = blockIdx.x*blockDim.x + threadIdx.x;
    unsigned int y = blockIdx.y*blockDim.y + threadIdx.y;
    unsigned int i = y*width+x;

    float2 slope;
    if ((x > 0) && (y > 0) && (x < width-1) && (y < height-1)) {
        slope.x = h[i+1] - h[i-1];
        slope.y = h[i+width] - h[i-width];
    } else {
        slope = make_float2(0.0f, 0.0f);
    }
    slopeOut[i] = slope;
}



///////////////////////////////////////////////////////////////////////////// 
/////////////////////////////// surface wakes /////////////////////////////// 
///////////////////////////////////////////////////////////////////////////// 

__device__ float fract(float x)
{
	return x - floor(x);
}

__device__ float* buf_ptr(const WakeBufInfo &buf, int x, int y)
{
	return (float*)((char*)buf.d_ptr + buf.pitch * y) + x;
}

__global__ void addWakeDropsKernel(int drop_n, const WakeDrop* d_drops, 
									float grid_step, WakeBufInfo buf)
{
    unsigned int i = blockIdx.x*blockDim.x + threadIdx.x;
	if (i >= drop_n)
		return;

	const int N = buf.size;
	float3 pos = d_drops[i].pos;
	float grid_extent = grid_step * N;
	int xi = fract(pos.x / grid_extent) * N;
	int yi = fract(pos.y / grid_extent) * N;
	if (xi >= N || yi >= N || xi < 0 || yi < 0)
		return;

	float ofs = d_drops[i].offset;
	float ofs2 = 0.5f * ofs;
	float ofs8 = 0.125f * ofs; 

	/*atomicAdd(buf_ptr(buf, xi, yi), ofs2);
	atomicAdd(buf_ptr(buf, (xi+1)%N, yi), ofs8);
	atomicAdd(buf_ptr(buf, (xi+N-1)%N, yi), ofs8);
	atomicAdd(buf_ptr(buf, xi, (yi+1)%N), ofs8);
	atomicAdd(buf_ptr(buf, xi, (yi+N-1)%N), ofs8);*/

	*buf_ptr(buf, xi, yi)			+= ofs2;
	/**buf_ptr(buf, (xi+1)%N, yi)		+= ofs8;
	*buf_ptr(buf, (xi+N-1)%N, yi)	+= ofs8;
	*buf_ptr(buf, xi, (yi+1)%N)		+= ofs8;
	*buf_ptr(buf, xi, (yi+N-1)%N)	+= ofs8;*/
}

texture<float, 2, cudaReadModeElementType> t1_tex;

__global__ void simualteWakeKernel(	float grid_step, WakeBufInfo dst_buf)
{
    unsigned int x = blockIdx.x*blockDim.x + threadIdx.x;
    unsigned int y = blockIdx.y*blockDim.y + threadIdx.y;

	int size = dst_buf.size;
	if (x >= size || y >= size)
		return;


	float *dst_ptr = (float*)((char*)dst_buf.d_ptr + dst_buf.pitch * y) + x;

	const float dt = 1.0 / 60.0;
	const float c = 0.6f;
	const float damping = 0.98f;

	float dp = 1.0f / dst_buf.size;
	float xf = (x + 0.5f)*dp, yf = (y + 0.5f)*dp;	

	float f = tex2D(t1_tex, xf, yf);
	float fn = tex2D(t1_tex, xf, yf+dp);
	float fs = tex2D(t1_tex, xf, yf-dp);
	float fe = tex2D(t1_tex, xf+dp, yf);
	float fw = tex2D(t1_tex, xf-dp, yf);
	float fp = *dst_ptr;


	float lap = (fn + fs + fe + fw - 4.0 * f) / (grid_step*grid_step);
	float new_f = f + damping * (f - fp) + dt*c*c*lap;

	*dst_ptr = new_f;
}



__global__ void addWakeToWave(WakeBufInfo src0, WakeBufInfo src1, float2* dst)
{
    unsigned int x = blockIdx.x*blockDim.x + threadIdx.x;
    unsigned int y = blockIdx.y*blockDim.y + threadIdx.y;

	int size = src0.size;
	if (x >= size || y >= size)
		return;
	
	float h0 = *buf_ptr(src0, x, y);
	float h1 = *buf_ptr(src1, x, y);
	
	dst[y*size + x].x += h1;
	dst[y*size + x].y = (h1 - h0) / 60.0; // DHG : 60 sec - Ok, greate HACK, to avoid foam on waves
}


extern "C" 
void cudaGenerateSpectrumKernel(float2* d_h0, float2 *d_ht, 
                                unsigned int width, unsigned int height, 
                                float animTime, float patchSize)
{
    dim3 block(8, 8, 1);
    dim3 grid(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);
    generateSpectrumKernel<<<grid, block>>>(d_h0, d_ht, width, height, animTime, patchSize);
}

extern "C"
void cudaCalculateSlopeKernel(  float* hptr, float2 *slopeOut, 
                                unsigned int width, unsigned int height)
{
    dim3 block(8, 8, 1);
    dim3 grid2(cuda_iDivUp(width, block.x), cuda_iDivUp(height, block.y), 1);
    calculateSlopeKernel<<<grid2, block>>>(hptr, slopeOut, width, height);
}

extern "C"
void cudaAddWakeDropsKernel( int drop_n, const WakeDrop* d_drops, const WakeBufInfo &buf, float grid_step)
{
	int block_size = 256;
	addWakeDropsKernel<<<cuda_iDivUp(drop_n, block_size), block_size>>>(drop_n, d_drops, grid_step, buf);
}

extern "C"
void cudaSimualteWakeKernel( float grid_step, const WakeBufInfo *bufs, int cur_buf)
{
	int size = bufs[0].size;

	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();
    
	dim3 block(8, 8, 1);
    dim3 grid(cuda_iDivUp(size, block.x), cuda_iDivUp(size, block.y), 1);

	const WakeBufInfo &dst_buf = bufs[cur_buf];
	const WakeBufInfo &t1_buf = bufs[1 - cur_buf];

	t1_tex.addressMode[0] = cudaAddressModeWrap; 
	t1_tex.addressMode[1] = cudaAddressModeWrap; 
	t1_tex.filterMode = cudaFilterModePoint; 
	t1_tex.normalized = 1;

	size_t ofs = 0;
	cudaBindTexture2D(&ofs, t1_tex, t1_buf.d_ptr, channelDesc, size, size, t1_buf.pitch);

	simualteWakeKernel<<<grid, block>>>(grid_step, dst_buf);
}

extern "C"
void cudaAddWakeToWave( const WakeBufInfo *bufs, int cur_buf, float2 *d_dst)
{
	int size = bufs[0].size;
    
	dim3 block(8, 8, 1);
    dim3 grid(cuda_iDivUp(size, block.x), cuda_iDivUp(size, block.y), 1);
	addWakeToWave<<<grid, block>>>(bufs[1-cur_buf], bufs[cur_buf], d_dst);
}
